import pickle
import pandas as pd
from surprise import Dataset, Reader

# 加载已保存的模型
def load_model(path):
    with open(path, 'rb') as f:
        return pickle.load(f)

# 生成Top 10推荐电影
def generate_top10_recommendations(user_id, model, movies_dict):
    predictions = [model.predict(user_id, movie_id) for movie_id in movies_dict.keys()]
    predictions.sort(key=lambda x: x.est, reverse=True)
    top10_movies = [pred.iid for pred in predictions[:10]]
    top10_titles = [movies_dict[movie_id] for movie_id in top10_movies]
    return top10_titles

def main():
    # 1. 加载训练好的模型
    try:
        svd_model = load_model('svd_recommender_model.pkl')
        knn_model = load_model('knn_recommender_model.pkl')
        print("模型加载成功！")
    except FileNotFoundError:
        print("错误：找不到模型文件，请确保 pkl 文件存在。")
        return

    # 2. 加载电影字典（movieId -> title）
    try:
        movies_df = pd.read_csv('processed_movies.csv')  # 确保你有一个movies.csv文件，包含movieId和title
        movies_dict = dict(zip(movies_df['movieId'], movies_df['title']))
    except FileNotFoundError:
        print("错误：未找到 movies.csv 文件，请提供包含 movieId 和 title 的电影数据。")
        return

    # 3. 设置要推荐的用户ID
    user_id = 1001  # 可以根据需要更换用户ID

    # 4. 生成推荐结果
    print(f"\n正在为用户 {user_id} 生成推荐...")

    top10_svd = generate_top10_recommendations(user_id, svd_model, movies_dict)
    top10_knn = generate_top10_recommendations(user_id, knn_model, movies_dict)

    print("\nSVD 推荐 Top 10:")
    for idx, title in enumerate(top10_svd, 1):
        print(f"{idx}. {title}")

    print("\nKNN 推荐 Top 10:")
    for idx, title in enumerate(top10_knn, 1):
        print(f"{idx}. {title}")

if __name__ == "__main__":
    main()
